import os
import requests
import json
import asyncio
import discord
import aiohttp
from depend.master import is_master

history_map = {}
prompt = "你是op要喝op果奶编写的y-tian-discord-bot,服务于阴天社区"
model_id = "gpt-4"
apikey = "nGwch_26Tg4U_yoL9szF_B3YgbW2I32nbKlg3vA32ag"


async def fetch_response_from_api(history, message_content):
  async with aiohttp.ClientSession() as session:
    data = {
      "model": {
        "id": model_id,
        "name": model_id,
        "maxLength": 12000,
        "tokenLimit": 4000
      },
      "messages": history,
      "key": apikey,
      "prompt": prompt,
      "temperature": 1
    }
    async with session.post("https://cgptf.vercel.app/api/chat",
                            headers={"Content-Type": "application/json"},
                            json=data) as response:
      return await response.text()


def fetch_model_list():
  response = requests.get("https://chimeragpt.adventblocks.cc/v1/models",
                          headers={"Authorization": "Bearer " + apikey})
  models = json.loads(response.text)["data"]
  return [(i + 1, model['id']) for i, model in enumerate(models)]


class yuanshen:

  def __init__(self, client):
    self.client = client
    self.models = []

  async def my_on_message(self, message):
    global history_map, prompt
    if message.author.bot:
      return False
    user_id = str(message.author.id)
    history = history_map.setdefault(user_id, [])

    if message.content.startswith('oo设定列表'):
      current_folder = os.path.dirname(os.path.abspath(__file__))
      prompts_folder = os.path.join(current_folder, 'prompts')
      prompt_files = [
        f[:-4] for f in os.listdir(prompts_folder) if f.endswith('.txt')
      ]
      if prompt_files:
        formatted_prompt_list = "\n".join(
          [f"{i + 1}. {prompt}" for i, prompt in enumerate(prompt_files)])
        await message.channel.send('Here are the available prompts:\n' +
                                   formatted_prompt_list)
      else:
        await message.channel.send('There are no prompts available.')
      return True

    if message.content.startswith('oo使用设定'):
      if is_master(message):
        chosen_prompt = message.content[len('oo使用设定'):].strip()
        current_folder = os.path.dirname(os.path.abspath(__file__))
        file_path = os.path.join(current_folder,
                                 f'prompts/{chosen_prompt}.txt')
        if os.path.isfile(file_path):
          with open(file_path, 'r') as file:
            prompt = file.read().strip()
          await message.channel.send('Your prompt has been updated.')
        else:
          await message.channel.send(
            '未找到该设定\n请发送：  oo添加设定  添加设定  （仅主人）\n或发送：  oo设定列表 查看已有设定')
      else:
        await message.channel.send("你没有权限执行此操作。")
      return True

    if message.content.startswith('oo结束对话'):
      history_map[user_id] = []
      await message.channel.send('Your chat history has been cleared.')
      return True

    if message.content.startswith('oo结束全部对话'):
      if is_master(message):
        history_map.clear()
        await message.channel.send(
          'Chat history has been cleared for all users.')
      else:
        await message.channel.send("你没有权限执行此操作。")
      return True

    if message.content.startswith('oo模型列表'):
      model_list = fetch_model_list()
      self.models = model_list
      formatted_model_list = "\n".join(f"{i}. {model}"
                                       for i, model in model_list)
      await message.channel.send('Here are the available models:\n' +
                                 formatted_model_list)
      return True

    if message.content.startswith('oo使用模型'):
      if is_master(message):
        model_index = int(message.content[len('oo使用模型'):].strip()) - 1
        if 0 <= model_index < len(self.models):
          chosen_model = self.models[model_index][1]
          global model_id
          model_id = chosen_model
          await message.channel.send(f'模型已设置为: {chosen_model}')
        else:
          await message.channel.send('请输入有效的模型序号')
      else:
        await message.channel.send("你没有权限执行此操作。")
      return True

    if message.content.startswith('oo查看设定'):
      chosen_prompt = message.content[len('oo查看设定'):].strip()
      current_folder = os.path.dirname(os.path.abspath(__file__))
      file_path = os.path.join(current_folder, f'prompts/{chosen_prompt}.txt')
      if os.path.isfile(file_path):
        with open(file_path, 'r') as file:
          prompt_content = file.read().strip()

        await message.channel.send('请选择发送模式：\nA：分段发送模式\nB：文件发送模式')

        def check_author_and_content(msg):
          return msg.author.id == message.author.id and msg.content in [
            'A', 'B'
          ]

        try:
          user_choice = await self.client.wait_for(
            'message', check=check_author_and_content, timeout=60)

          if user_choice.content == 'A':
            # 分段发送模式，以代码块形式发送
            for i in range(0, len(prompt_content), 2000 - 6):
              formatted_content = f"```{prompt_content[i:i + 2000 - 6]}```"
              await message.channel.send(formatted_content)
          elif user_choice.content == 'B':
            # 文件发送模式
            with open(file_path, 'rb') as file:
              await message.channel.send(f'设定 {chosen_prompt} 的内容：',
                                         file=discord.File(
                                           file, f'{chosen_prompt}.txt'))
        except asyncio.TimeoutError:
          await message.channel.send('超时，请重新发送命令。')
      else:
        await message.channel.send(
          '未找到该设定\n请发送：  oo添加设定  添加设定  （仅主人）\n或发送：  oo设定列表 查看已有设定')
      return True

    if message.content.startswith('oo'):
      chat_content = message.content[len('oo'):].strip()
      history.append({"role": "user", "content": chat_content})
      response = await fetch_response_from_api(history, chat_content)
      history.append({"role": "assistant", "content": response})
      await message.channel.send(response)
      return True

    return False
